package com.tsj.message;

import com.tsj.message.config.KafkaProperties;
import org.apache.kafka.clients.producer.*;
import org.apache.kafka.common.serialization.StringSerializer;

import java.util.Properties;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.Future;

/**
 * 生产者发送确认一批后，再发送下一批
 * 启动参数：sync
 */
public class Producer extends Thread {
    private final KafkaProducer<String, String> producer;
    private final String topic;
    private final Boolean isAsync;
    private int numRecords;
    private final CountDownLatch latch;

    private String keyPrefix;

    public Producer(final String client_id,
                    final String topic,
                    final Boolean isAsync,
                    final String transactionalId,
                    final boolean enableIdempotency,
                    final int numRecords,
                    final int transactionTimeoutMs,
                    final CountDownLatch latch,
                    final String keyPrefix) {
        Properties props = new Properties();
        props.put(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG, KafkaProperties.KAFKA_SERVER_URL + ":" + KafkaProperties.KAFKA_SERVER_PORT);
        props.put(ProducerConfig.CLIENT_ID_CONFIG, client_id);
        props.put(ProducerConfig.KEY_SERIALIZER_CLASS_CONFIG, StringSerializer.class.getName());
        props.put(ProducerConfig.VALUE_SERIALIZER_CLASS_CONFIG, StringSerializer.class.getName());
        if (transactionTimeoutMs > 0) {
            props.put(ProducerConfig.TRANSACTION_TIMEOUT_CONFIG, transactionTimeoutMs);
        }
        if (transactionalId != null) {
            props.put(ProducerConfig.TRANSACTIONAL_ID_CONFIG, transactionalId);
        }
        props.put(ProducerConfig.ENABLE_IDEMPOTENCE_CONFIG, enableIdempotency);
//        props.put(ProducerConfig.ACKS_CONFIG, "0");

        producer = new KafkaProducer<>(props);
        this.topic = topic;
        this.isAsync = isAsync;
        this.numRecords = numRecords;
        this.latch = latch;
        this.keyPrefix = keyPrefix;
    }

    KafkaProducer<String, String> get() {
        return producer;
    }

    @Override
    public void run() {
        KeyGenerator keyGenerator = new KeyGenerator(keyPrefix, true);
        int remainNum = numRecords;
        while (remainNum > 0) {
            String messageKey = keyGenerator.nextKey();
            if (isAsync) { // Send asynchronously
                producer.send(new ProducerRecord<>(topic, messageKey, messageKey),
                        new DemoCallBack(System.currentTimeMillis(), messageKey, messageKey));
            } else { // Send synchronously
                try {
                    Future<RecordMetadata> send = producer.send(
                            new ProducerRecord<>(topic, messageKey, messageKey));
                    send.get();
                    System.out.println("Sent message: (" + messageKey + ", " + messageKey + ")");
                } catch (InterruptedException | ExecutionException e) {
                    e.printStackTrace();
                }
            }
            remainNum -= 1;
        }
        System.out.println("Producer sent " + numRecords + " records successfully");
        latch.countDown();

    }

    class KeyGenerator {
        private Integer cur;
        private Integer step;
        private String keyPrefix;
        private boolean immutableKey;

        public KeyGenerator(String keyPrefix, boolean immutableKey, Integer start, Integer step) {
            this.step = step;
            this.cur = start;
            this.immutableKey = immutableKey;
            this.keyPrefix = keyPrefix;
        }

        public KeyGenerator(String keyPrefix, boolean immutableKey) {
            this.immutableKey = immutableKey;
            this.keyPrefix = keyPrefix;
        }

        public String nextKey() {
            if (immutableKey) {
                return keyPrefix;
            }
            String messageStr = keyPrefix + "_" + cur;
            cur += step;
            return messageStr;
        }
    }
}

class DemoCallBack implements Callback {

    private final long startTime;
    private final String key;
    private final String message;

    public DemoCallBack(long startTime, String key, String message) {
        this.startTime = startTime;
        this.key = key;
        this.message = message;
    }

    /**
     * A callback method the user can implement to provide asynchronous handling of request completion. This method will
     * be called when the record sent to the server has been acknowledged. When exception is not null in the callback,
     * metadata will contain the special -1 value for all fields except for topicPartition, which will be valid.
     *
     * @param metadata  The metadata for the record that was sent (i.e. the partition and offset). An empty metadata
     *                  with -1 value for all fields except for topicPartition will be returned if an error occurred.
     * @param exception The exception thrown during processing of this record. Null if no error occurred.
     */
    public void onCompletion(RecordMetadata metadata, Exception exception) {
        long elapsedTime = System.currentTimeMillis() - startTime;
        if (metadata != null) {
            System.out.println(
                    "message(" + key + ", " + message + ") sent to partition(" + metadata.partition() +
                            "), " +
                            "offset(" + metadata.offset() + ") in " + elapsedTime + " ms");
        } else {
            System.out.println("send metadata offset:" + metadata.offset());
            exception.printStackTrace();
        }
    }
}